import matplotlib.pyplot as plt
from collections import defaultdict
import argparse
import numpy as np

# 从文件中读取并解析区间数据的函数
def read_kmer_data(file_path):
    interval_data = defaultdict(dict)
    current_interval = None

    with open(file_path, 'r') as file:
        for line in file:
            line = line.strip()
            if line.startswith("Interval:"):
                current_interval = line.split("Interval: ")[1].strip()
                print(f"Found interval: {current_interval}")  # 调试信息
            elif line and current_interval:
                parts = line.split(":")
                if len(parts) == 2:
                    kmer = parts[0].strip()
                    try:
                        count = int(parts[1].strip())
                        # 仅处理kmer长度等于5的数据
                        if len(kmer) == 5:
                            interval_data[current_interval][kmer] = count
                            print(f"Added kmer: {kmer}, count: {count} to interval: {current_interval}")  # 调试信息
                    except ValueError:
                        pass

    return interval_data

# 计算每个kmer在不同区间中的变化比例
def calculate_kmer_changes(interval_data):
    if '[0,0]' not in interval_data:
        print("Error: Reference interval '[0,0]' not found in data.")  # 调试信息
        return {}, []

    base_counts = interval_data['[0,0]']
    base_total = sum(base_counts.values())
    print(f"Base total count for [0,0]: {base_total}")  # 调试信息

    # 计算初始比例
    base_ratios = {kmer: count / base_total for kmer, count in base_counts.items()}

    # 按顺序计算每个区间中的kmer比例变化
    interval_keys = sorted(interval_data.keys(), key=parse_interval_key)
    print(f"Intervals sorted: {interval_keys}")  # 调试信息

    kmer_changes = defaultdict(list)

    for interval in interval_keys:
        if interval == '[0,0]':
            continue

        interval_counts = interval_data[interval]
        interval_total = sum(interval_counts.values())
        print(f"Interval: {interval}, Total count: {interval_total}")  # 调试信息

        for kmer in base_ratios:
            base_ratio = base_ratios.get(kmer, 0)
            interval_ratio = interval_counts.get(kmer, 0) / interval_total if interval_total > 0 else 0
            change = abs(interval_ratio - base_ratio)  # 变化绝对值
            kmer_changes[kmer].append(change)
            print(f"K-mer: {kmer}, Change in interval {interval}: {change}")  # 调试信息

    return kmer_changes, interval_keys

# 对区间进行排序的函数
def parse_interval_key(interval):
    if interval == '[0,0]':
        return -1
    elif interval == '[100,100]':
        return 101
    else:
        # 提取区间中的数值部分，例如 (10,20] -> 10
        return int(interval.split(',')[0].replace('(', '').replace('[', '').strip())

# 找出方差最大和最小的top5 kmer
def find_top_kmers_by_variance(kmer_changes, top_n=5):
    kmer_variances = {kmer: np.var(changes) for kmer, changes in kmer_changes.items()}
    print(f"Variance calculated for all kmers.")  # 调试信息

    # 找出方差最大的top_n个kmer
    top_max_kmers = sorted(kmer_variances.items(), key=lambda x: x[1], reverse=True)[:top_n]
    # 找出方差最小的top_n个kmer
    top_min_kmers = sorted(kmer_variances.items(), key=lambda x: x[1])[:top_n]

    print(f"Top {top_n} max variance kmers: {top_max_kmers}")  # 调试信息
    print(f"Top {top_n} min variance kmers: {top_min_kmers}")  # 调试信息

    return [kmer for kmer, _ in top_max_kmers], [kmer for kmer, _ in top_min_kmers]

# 绘制kmer的变化曲线
def plot_kmer_changes(kmer_changes, interval_keys, top_max_kmers, top_min_kmers):
    if not interval_keys:
        print("No interval keys available for plotting.")  # 调试信息
        return

    plt.figure(figsize=(14, 8))

    # 绘制变化最大的top5 kmer
    for kmer in top_max_kmers:
        if kmer in kmer_changes:
            plt.plot(interval_keys[1:], kmer_changes[kmer], label=f'Max Variance: {kmer}', linestyle='-', marker='o')

    # 绘制变化最小的top5 kmer
    for kmer in top_min_kmers:
        if kmer in kmer_changes:
            plt.plot(interval_keys[1:], kmer_changes[kmer], label=f'Min Variance: {kmer}', linestyle='--', marker='x')

    plt.xlabel('Methylation Interval')
    plt.ylabel('Change Proportion (0% to 3.5%)')
    plt.title('Top 5 K-mers with Maximum and Minimum Variance in Changes Across Intervals')
    plt.xticks(rotation=45)
    plt.ylim(0, 0.035)  # 纵轴范围为 0% 到 10%
    plt.legend()
    plt.tight_layout()
    plt.show()

# 主函数，用于整体控制逻辑
def main():
    # 设置命令行参数解析
    parser = argparse.ArgumentParser(description="读取kmer统计文件并绘制方差最大的和最小的kmer变化图")
    parser.add_argument("file_path", help="输入的kmer统计文件路径")
    args = parser.parse_args()

    # 读取kmer数据
    interval_data = read_kmer_data(args.file_path)

    # 计算每个kmer在不同区间的变化比例
    kmer_changes, interval_keys = calculate_kmer_changes(interval_data)

    if not kmer_changes:
        print("No kmer changes data found. Exiting.")  # 调试信息
        return

    # 找出方差最大的top5和最小的top5的kmer
    top_max_kmers, top_min_kmers = find_top_kmers_by_variance(kmer_changes)

    # 绘制这些kmer的变化曲线
    plot_kmer_changes(kmer_changes, interval_keys, top_max_kmers, top_min_kmers)

if __name__ == "__main__":
    main()
